/**
 * @file AMGPCGsolver.h
 * @author shanyoumu
 * @brief 鉴于eigen的矩阵速度太慢，现将矩阵改成dealii并尝试对接dealii的PCG求解器
 * @version 0.1
 * @date 2023-03-11
 *
 * @copyright Copyright (c) 2023
 *
 */

#ifndef _AMGPCGSOLVER_H_
#define _AMGPCGSOLVER_H_

#include <algorithm>
#include <chrono>
#include <cmath>
#include <fstream>
#include <iostream>
#include <list>
#include <vector>

#include <deal.II/lac/sparse_direct.h>
#include <deal.II/lac/sparse_matrix.h>
#include <deal.II/lac/sparsity_pattern.h>
#include <deal.II/lac/vector.h>

using DealiiSpMat = dealii::SparseMatrix<double>;
using DealiiVectorXd = dealii::Vector<double>;
using dealii::SparseDirectUMFPACK;
using dealii::SparsityPattern;
using std::cout;
using std::endl;
using std::fabs;
using std::list;
using std::vector;
using u32 = unsigned int;

/**
 * @brief 建立粗网格的组件
 *
 */
class BuildMesh;
/**
 * @brief 代数多重网格v-循环
 *
 */
class AMG;
/**
 * @brief 高斯塞尔德迭代, Mx=r迭代s次
 *
 * @param M
 * @param x
 * @param rhs
 * @param itnum
 */
void Gauss_Seidel(const DealiiSpMat &M, DealiiVectorXd &x,
                  const DealiiVectorXd &rhs, const int s);
double L2_Error(DealiiVectorXd &v1, DealiiVectorXd &v2);

class BuildMesh {
  private:
    DealiiSpMat *M;
    double theta;
    SparsityPattern SpItp;
    DealiiSpMat ItplatMat; //插值算子
    vector<u32> Cpoints;   //保存选取的粗网格点
    vector<u32> rowlength; //这两个用来保存每一行/列非零元的个数
    vector<u32> collength;

  public:
    BuildMesh(DealiiSpMat &Spt, double _theta = 0.25) {
        M = &Spt;
        theta = _theta;
    };
    void build();                                         //建立插值矩阵
    DealiiSpMat &GetMatP() { return ItplatMat; };         //获得插值矩阵
    const vector<u32> &GetCorsePnt() { return Cpoints; }; //获得粗网格点集
    vector<u32> &Getlength() { return rowlength; }; //获得每行非零元个数
};

class AMG {
  private:
    vector<SparsityPattern> Vec_SPM;   //保存每一层网格矩阵的pattern
    vector<SparsityPattern> Vec_SPP;   //保存投影算子的pattern
    vector<SparsityPattern> Vec_SPPt;  //保存插值算子的pattern
    vector<DealiiSpMat> Vec_Mat;       //保存每一层网格的矩阵
    vector<DealiiSpMat> Interpolation; //保存每一层网格的插值算子
    vector<DealiiSpMat> Restrict;      //保存每一层网格的投影算子
    vector<DealiiVectorXd> F_Rhs;      //保存右端项
    vector<DealiiVectorXd> U_Lhs;      //保存每一层的解
    SparseDirectUMFPACK A_direct;      //最粗网格采用直接法求解

    int MaxLayerNum;
    int smooth_num;
    int N_layer;      //网格层数
    double Tolerance; //迭代精度
    double Theta;     //强连接点阈值
  public:
    AMG() = default;
    /// @brief 建立AMG
    /// @param _Spt 求解的矩阵
    /// @param _Tol 判断迭代收敛的阈值
    /// @param _theta 强连接点的阈值, 一般取0.25
    /// @param smooth_num 每一层矩阵做GS迭代的次数
    /// @param _maxlayernum 最大网格层数
    AMG(DealiiSpMat &_Spt, double _Tol = 1e-13, double _theta = 0.25,
        int smooth_num = 2, int _maxlayernum = 10);
    DealiiVectorXd &solve(const DealiiVectorXd &Rhs); //用AMG求解方程
    void Vcycle(int _layerid); //从第_layerid层开始做vcycle
    void Vcycle();             //做一次完整的vcycle
    DealiiVectorXd &
    Vcycle(const DealiiVectorXd &Rhs); //做一次vcycle后得到的近似解
    void Vcycleonce(DealiiVectorXd &x,
                    const DealiiVectorXd &Rhs); //对x做一次vcycle
    DealiiVectorXd &Vcycleonce(DealiiVectorXd &Rhs) {
        F_Rhs[0] = Rhs;
        U_Lhs[0].reinit(Rhs.size());
        Vcycle(0);
        return U_Lhs[0];
    }
    DealiiSpMat &GetlastMat() { return Vec_Mat.back(); };
};

void RebuildMat(DealiiSpMat &M); //将矩阵转为M-阵

/// PCG部分

class AMGPCG {
  private:
    double Tolerance;
    int iteratornum = 0;
    int Max_itenum; //设置最大迭代次数
    int Smoothstep;

    bool rebuild;
    double theta;        // AMG强连接点的阈值
    DealiiSpMat A;       //求解的矩阵
    SparsityPattern SpA; //求解的矩阵pattern
    DealiiVectorXd Rhs;  //方程的右端项
    AMG *AmgGrid; //保存AMG的每一层网格的矩阵，重复利用以求解

    /// PCG迭代用到的变量
    DealiiVectorXd zk;
    DealiiVectorXd zk_1;

    DealiiVectorXd rk;
    DealiiVectorXd rk_1;
    DealiiVectorXd pk;
    DealiiVectorXd pk_1;
    DealiiVectorXd xk;
    double alphak;
    double betak;

  public:
    AMGPCG(DealiiSpMat &_Mat, bool _rebuild = true, int _maxitnum = 200,
           double _tol = 1e-13, double theta = 0.25, int _smoothstep = 1)
        : Max_itenum(_maxitnum), rebuild(_rebuild), Tolerance(_tol),
          theta(theta), Smoothstep(_smoothstep) {

        const SparsityPattern& spM = _Mat.get_sparsity_pattern();
        SpA.copy_from(spM);
        A.reinit(SpA);
        for(u32 i = 0; i < _Mat.m(); i++) {
          for(auto it = _Mat.begin(i); it != _Mat.end(i); it++) {
            A.add(i,it->column(),it->value());
          }
        }
        // A.copy_from(_Mat);
        // const SparsityPattern& spM = _Mat.get_sparsity_pattern();
        // const std::size_t * rowstart = spM.get_rowstart_indices();
        // const u_int * colnums = spM.get_column_numbers();
        //
        // SpA.copy_from(spM);
        // A.reinit(SpA);
        // for (u_int j = 0;j < _Mat.m();j ++) {
        //     for (u_int k = rowstart[j] ;k < rowstart[j+1]; k++) {
        //     A.add(j, colnums[k], _Mat.global_entry(k));
        //     }
        // }
        BuildAMG();
    };
    ~AMGPCG(){ delete AmgGrid;};
    DealiiVectorXd solve(DealiiVectorXd& _Rhs);           //用AMGPCG求解方程
    void BuildAMG();
    int GetItNum() {return iteratornum;};
};
#endif
